#include <sys/epoll.h>
#include <errno.h>

#include "asyncio.hpp"

namespace asyncio {

    template <typename T, typename Fn>
    T schedule_fd_async(coro::Context<IoResult> &ctx, int fd, int event, Fn fn)
    {
        coro::Waker<IoResult> *waker = new coro::Waker<IoResult>(ctx.waker());
        int r = ctx.get_executor()->io_schedule(fd, event | EPOLLRDHUP,
            [waker](int evt) {
                IoResult res = {
                    .events = evt,
                };
                waker->wake(res);
                delete waker;
            });
        if (r < 0) {
            errno = -r;
            return -1;
        }

        std::optional<IoResult> res = ctx.yield();
        if (res) {
            auto _res = res.value();
            if (_res.events & (EPOLLHUP | EPOLLERR | EPOLLRDHUP)) {
                errno = ENOENT;
                return -1;
            } else if (_res.events & event) {
                return fn();
            } else {
                errno = EINVAL;
                return -1;
            }
        } else {
            throw std::runtime_error("failed to schedule fd");
        }
    }

    int accept_async(coro::Context<IoResult> &ctx,
                     int sockfd, struct sockaddr *addr, socklen_t *addrlen)
    {
        return schedule_fd_async<int>(ctx, sockfd, EPOLLIN,
            [sockfd, addr, addrlen]() -> int {
                return accept4(sockfd, addr, addrlen, SOCK_NONBLOCK | SOCK_CLOEXEC);
            });
    }

    int connect_async(coro::Context<IoResult> &ctx,
                      int sockfd, const struct sockaddr *addr, socklen_t addrlen)
    {
        // Start connect and wait for the socket to be writable
        int r = connect(sockfd, addr, addrlen);
        if (r < 0 && errno != EINPROGRESS) {
            return -errno;
        }
        return schedule_fd_async<int>(ctx, sockfd, EPOLLOUT,
            []() -> int {
                return 0;
            });
    }

    ssize_t read_async(coro::Context<IoResult> &ctx,
                       int fd, void *buf, size_t count)
    {
        return schedule_fd_async<ssize_t>(ctx, fd, EPOLLIN,
            [fd, buf, count]() -> ssize_t {
                return read(fd, buf, count);
            });
    }

    ssize_t write_async(coro::Context<IoResult> &ctx,
                        int fd, const void *buf, size_t count)
    {
        return schedule_fd_async<ssize_t>(ctx, fd, EPOLLOUT,
            [fd, buf, count]() -> ssize_t {
                return write(fd, buf, count);
            });
    }
};
